#include "CommonGPU.h"
#include <cuda_runtime_api.h>
#include <stdio.h>
#include <stdlib.h>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/scan.h>
#include <iostream>
#include <thrust/remove.h>
#include <device_launch_parameters.h>
#include <cub/cub.cuh>
using namespace std;

__device__ void blk_mergesort_shared_memory(int* shared_array, int* temp_array, int size){
	int rounds = ceil(log2((double)size));

	int base = 1;
	for(int i = 1; i <= rounds; i++){
		
		if(threadIdx.x % (2*base) == 0){
			int start_idx_x = threadIdx.x;
			int end_idx_x = start_idx_x + base;
			end_idx_x = end_idx_x < size ? end_idx_x : size;

			int start_idx_y = end_idx_x;
			int end_idx_y = start_idx_y + base;
			end_idx_y = end_idx_y < size ? end_idx_y : size;

			int x_pointer = start_idx_x;
			int y_pointer = start_idx_y;

			int output_ptr = x_pointer;
			while(x_pointer < end_idx_x || y_pointer < end_idx_y){
				if(x_pointer >= end_idx_x)
					temp_array[output_ptr++] = shared_array[y_pointer++];
				else if(y_pointer >= end_idx_y)
					temp_array[output_ptr++] = shared_array[x_pointer++];
				else if(shared_array[x_pointer] < shared_array[y_pointer])
					temp_array[output_ptr++] = shared_array[y_pointer++];
				else
					temp_array[output_ptr++] = shared_array[x_pointer++];
			}			
		}
		__syncthreads();
		if(threadIdx.x < size){
			shared_array[threadIdx.x] = temp_array[threadIdx.x];
		}
		__syncthreads();
		base *= 2;
	}
}

// data transpose, interleaved to consecutive, 4 elements assumption with blockDim.x = SORTED_SHARED_SIZE
__device__ void interleave_to_consecutive(int* local_data, bool to_print){
	__shared__ int shared_data[4*SORTED_SHARED_SIZE];
	for(int i = 0; i < 4; i++) shared_data[i * blockDim.x + threadIdx.x] = local_data[i];
	__syncthreads();

	if(threadIdx.x == 0 && blockIdx.x == 1 && to_print){
		for(int i = 0; i < 4; i++){
			printf("++++++++++++++++++++++++++++++++\n");
			for(int j = 0; j < SORTED_SHARED_SIZE; j++){
				printf("%d ",shared_data[i*SORTED_SHARED_SIZE+j]);
			}
			printf("\n");
		}
		printf("end of the share structure interleave_to_consecutive\n\n");

	}
	__syncthreads();

	for(int i = 0; i < 4; i++) local_data[i] = shared_data[threadIdx.x * 4 + i];
}
// data transpose, consecutive to interleaved, 4 elements assumption with blockDim.x = SORTED_SHARED_SIZE
__device__ void consecutive_to_interleave(int* local_data,bool to_print){
	__shared__ int shared_data[4*SORTED_SHARED_SIZE];
	for(int i = 0; i < 4; i++) shared_data[threadIdx.x * 4 + i] = local_data[i];
	__syncthreads();

	if(threadIdx.x == 0 && blockIdx.x == 1 && to_print){
		for(int i = 0; i < 4; i++){
			printf("++++++++++++++++++++++++++++++++\n");
			for(int j = 0; j < SORTED_SHARED_SIZE; j++){
				printf("%d ",shared_data[i*SORTED_SHARED_SIZE+j]);
			}
			printf("\n");
		}
		printf("end of the share structure consecutive_to_interleave\n\n");

	}
	__syncthreads();

	for(int i = 0; i < 4; i++) local_data[i] = shared_data[i * blockDim.x + threadIdx.x];
}

__global__ void sort_within_blks_tuned(int* hash_table, int* hash_index, int hash_index_offset,int* sorted_hash_table, int num_rounds, int MAX_BITS){

	// SORTED_SHARED_SIZE = blockDim.x, must be garranteed
	
	int index = blockIdx.x*hash_index_offset-1;
	int blk_start = (blockIdx.x == 0) ? 0 : hash_index[index];
	int blk_end = hash_index[index+hash_index_offset];
	int blk_size = blk_end - blk_start;

	int round = ceil((double)blk_size / SORTED_SHARED_SIZE);
	int bit_round = ceil((double)MAX_BITS / 2); // default use 4 as radix

	if(threadIdx.x == 0) printf("block %d start at %d and end at %d\n",blockIdx.x,blk_start,blk_end);

	for(int bit_shift = 0; bit_shift < bit_round; bit_shift++){

		//if(blockIdx.x == 1 && threadIdx.x == 0) printf("bit_shift %d with blk_size %d in bit_shift %d ----- \n ",bit_shift,blk_size,bit_shift);
		//__syncthreads();

		int total_pos[4];
		for(int i = 0; i < 4; i++) total_pos[i] = 0;

		for(int i = 0; i < round; i++){
			int idx = blk_start+i*SORTED_SHARED_SIZE+threadIdx.x;

			// load data
			unsigned int data = (idx < blk_end) ? hash_table[idx] : 0;
			int bucket = (data >> (2*bit_shift)) % 4;

			int pos[4];
			for(int i = 0; i < 4; i++) pos[i] = 0;
			pos[bucket] = (idx < blk_end);

			// prefix count
			{
				interleave_to_consecutive(pos,false);

				typedef cub::BlockScan<int,SORTED_SHARED_SIZE> BlockScan;
				__shared__ typename BlockScan::TempStorage temp_storage;
				BlockScan(temp_storage).ExclusiveSum(pos, pos);
				__syncthreads();

				consecutive_to_interleave(pos,false);
			}

			__shared__ int partitioned_array[SORTED_SHARED_SIZE];
			if(idx < blk_end){
				partitioned_array[pos[bucket]] = data;
			}
			__syncthreads();

			if(blockIdx.x == 1 && threadIdx.x == 0){
				for(int k = 0; k < SORTED_SHARED_SIZE; k++) printf("%d ",partitioned_array[k]);
				printf("\n");
			}
			__syncthreads();
			if(blockIdx.x == 1){
				printf("thread %d has data %d and bucket %d and bucket_test %d: \n",threadIdx.x,data,bucket,(data/(4*bit_shift))%4);
			}
			__syncthreads();

			// update global count in total_pos array
			data = partitioned_array[threadIdx.x];
			bucket = (data >> (2*bit_shift)) % 4;
			total_pos[bucket] = total_pos[bucket] + (int)(idx < blk_end);

			if(idx < blk_end){
				hash_table[idx] = data;
			}
			__syncthreads();
		}

		// global prefix count
		{
			interleave_to_consecutive(total_pos,true);

			typedef cub::BlockScan<int,SORTED_SHARED_SIZE> BlockScan;
			__shared__ typename BlockScan::TempStorage temp_storage;
			BlockScan(temp_storage).ExclusiveSum(total_pos, total_pos);
			__syncthreads();

			consecutive_to_interleave(total_pos,true);
		}

		// output array
		for(int i = 0; i < round; i++){
			int idx = blk_start+i*SORTED_SHARED_SIZE+threadIdx.x;

			// load data
			unsigned int data = (idx < blk_end) ? hash_table[idx] : 0;
			int bucket = (data >> (2*bit_shift)) % 4;

			if(idx < blk_end){
				int pos = blk_start+total_pos[bucket];
				if(blockIdx.x == 1) printf("thread %d has data %d with bucket %d at bit_shift %d output at %d\n",
					threadIdx.x,data,bucket,bit_shift,pos);
				sorted_hash_table[blk_start+total_pos[bucket]] = data;
				total_pos[bucket]++;
			}
		}
		__syncthreads();

		for(int i = 0; i < round; i++){
			int idx = blk_start+i*SORTED_SHARED_SIZE+threadIdx.x;
			if(idx < blk_end){
				hash_table[idx] = sorted_hash_table[idx];
			}
		}
		__syncthreads();

	}

}

__global__ void sort_within_blks(int* hash_table, int* hash_index, int hash_index_offset,int* sorted_hash_table, int num_rounds){
	__shared__ int shared_array[SORTED_SHARED_SIZE];
	__shared__ int temp_array[SORTED_SHARED_SIZE];
	__shared__ int prefix_count[SORTED_SHARED_SIZE*4];

	int bid = blockIdx.x;
	int tid = threadIdx.x;

	int index = bid*hash_index_offset-1;
	int blk_start = (bid == 0) ? 0 : hash_index[index];
	int blk_end = hash_index[index+hash_index_offset];
	int blk_size = blk_end - blk_start;
	
	int round = ceil((double)blk_size / SORTED_SHARED_SIZE);


	if(tid == 0) printf("blk size is %d for block %d and round %d\n",blk_size,bid,round);
	//if(bid == 0 && tid == 0) printf("blk size is %d with round %d\n",blk_size,round);

	//if(tid == 0 && bid == 0) printf("round is %d blk_size %d blockDim %d\n",round,blk_size,blockDim.x);

	for(int i = 0; i < round; i++){
		int idx = blk_start+i*SORTED_SHARED_SIZE+tid;
		if(idx < blk_end){
			shared_array[tid] = hash_table[idx];
		}
		int offset = blk_size % SORTED_SHARED_SIZE;
		int sorted_size = (i == round -1 && offset != 0) ? offset : SORTED_SHARED_SIZE;
		__syncthreads();

		blk_mergesort_shared_memory(shared_array,temp_array,sorted_size);

		if(idx < blk_end){
			hash_table[idx] = shared_array[tid];
		}
		__syncthreads();
	}

	// if(round <= 1 || blk_size == 0) return;

	//// merge the sorted array
	//int big_rounds = ceil(log2((double)round));
	//int base = 1;
	//int* src = sorted_hash_table;
	//int* des = hash_table;

	//for(int i = 0; i < big_rounds; i++){
	//	int* temp = src;
	//	src = des;
	//	des = temp;

	//	int sub_iter = ceil((double)round / SORTED_SHARED_SIZE);

	//	for(int j = 0; j < sub_iter; j++){
	//		int tid_inner = j * SORTED_SHARED_SIZE + tid;

	//		if(tid_inner % (2*base) == 0){
	//			

	//			int start_idx_x = blk_start + tid_inner * SORTED_SHARED_SIZE;
	//			int end_idx_x = start_idx_x + base * SORTED_SHARED_SIZE;
	//			end_idx_x = min(end_idx_x,blk_end);

	//			int start_idx_y = end_idx_x;
	//			int end_idx_y = end_idx_x + base * SORTED_SHARED_SIZE;
	//			end_idx_y = min(end_idx_y,blk_end);

	//			int x_pointer = start_idx_x;
	//			int y_pointer = start_idx_y;

	//			int output_ptr = x_pointer;

	//			//if(bid == 1) printf("big round %d thread %d start_x %d start_y %d \n",i,tid,start_idx_x,end_idx_x);

	//			while(x_pointer < end_idx_x || y_pointer < end_idx_y){
	//				if(x_pointer >= end_idx_x)
	//					des[output_ptr++] = src[y_pointer++];
	//				else if(y_pointer >= end_idx_y)
	//					des[output_ptr++] = src[x_pointer++];
	//				else if(src[x_pointer] < src[y_pointer])
	//					des[output_ptr++] = src[y_pointer++];
	//				else
	//					des[output_ptr++] = src[x_pointer++];
	//			}
	//		}
	//	}

	//	__syncthreads();
	//	round = ceil((double)round /2);
	//	base *= 2;
	//}

	//if(big_rounds % 2 == 1) return;

	round = ceil((double)blk_size / SORTED_SHARED_SIZE);
	for(int i = 0; i < round; i++){
		int idx = blk_start+i*SORTED_SHARED_SIZE+tid;
		if(idx < blk_end){
			sorted_hash_table[idx] = hash_table[idx];
		}
	}

}

// output_index has number of element: number of blks * number of MAX_HASH_KEY
// first (number of threads) elements denote the index for hash key 0
// here number of threads per blk must be equal to MAX_HASH_KEY
__global__ void hash_partition(int* data, int data_size, int num_of_blks, int* output_index, int num_rounds){
	

	

	int tid = threadIdx.x + blockIdx.x * blockDim.x;

	//if(tid == 0){
	//	printf("data size xxxxxxxxxxxxxxxxx %d \n",data_size);
	//	printf("is enter loop: %d\n",tid < data_size);
	//}



	if(num_rounds != 3){ 

	__shared__ int hash_count[MAX_HASH_KEY];

	hash_count[threadIdx.x] = 0;
	__syncthreads();

	if(tid < data_size){
		int value = data[tid];
		int hash_value = (value * 2654435761) % MAX_HASH_KEY;
		//if(num_rounds == 3 && blockIdx.x % 128 == 0){
		//	printf("hash_value is %d\n",hash_value);
		//}
		atomicAdd(&hash_count[hash_value],1);
	}

	__syncthreads();

	//if(num_rounds == 3 && tid == 0){
	//	printf("shared hash count\n");
	//	for(int i = 0; i < MAX_HASH_KEY; i++){
	//		printf("%d ",hash_count[i]);
	//	}
	//}

	output_index[threadIdx.x*num_of_blks+blockIdx.x] = hash_count[threadIdx.x];
	}
}

__global__ void output_hash_table(int* data, int data_size, int num_of_blks,int* output_index, int* output_table){
	__shared__ int hash_count[MAX_HASH_KEY];

	hash_count[threadIdx.x] = (threadIdx.x == 0 && blockIdx.x == 0) ? 0 : output_index[threadIdx.x*num_of_blks+blockIdx.x-1];
	__syncthreads();


	int tid = threadIdx.x + blockIdx.x * blockDim.x;

	if(tid < data_size){
		int value = data[tid];
		int hash_value = (value * 2654435761) % MAX_HASH_KEY;
		int output_pos = atomicAdd(&hash_count[hash_value],1);
		output_table[output_pos] = value;
	}
}

__global__ void mark_duplicate(int* data, int data_size,int* remove_indicator){
	int tid = threadIdx.x + blockIdx.x * blockDim.x;

	if(tid < data_size){
		int value = data[tid];
		int s_value = (tid == 0) ? -1 : data[tid-1];
		remove_indicator[tid] = value == s_value;
	}
}

// return the new data size
int remove_duplicate(thrust::device_vector<int>& dvec,int data_size,int num_rounds){

	if(data_size <= 1) return data_size;
	//thrust::host_vector<int> hvec(dvec.begin(),dvec.begin()+data_size);
	//vector<vector<int> > hash_table(MAX_HASH_KEY);
	//for(int i = 0; i < data_size; i++){
	//	int value = hvec[i];
	//	int hash_value = (value * 2654435761) % MAX_HASH_KEY;
	//	hash_table[hash_value].push_back(value);
	//}
	//for(int i = 0; i < MAX_HASH_KEY; i++){
	//	cout << "value " << i << " size " << hash_table[i].size() << ":";
	//	for(int j = 0; j < hash_table[i].size(); j++){
	//		int value = hash_table[i][j];
	//		cout << "(" << value << "," << i << ") | ";
	//	}
	//	cout << endl;
	//}
	//cout << "start to compute" << endl;
	vector<float> exec_time;

	cudaEvent_t start, stop;
	float elapsedTime;
	cudaEventCreate(&start);
	cudaEventCreate(&stop);

	int num_of_blks = (data_size / MAX_HASH_KEY) + (data_size % MAX_HASH_KEY != 0);
	thrust::device_vector<int> ind_vec(num_of_blks * MAX_HASH_KEY);

	cudaEventRecord(start, 0);

	//cout << "hash partition with num_of_blks " << num_of_blks << endl;
	hash_partition<<<num_of_blks,MAX_HASH_KEY>>>(
		thrust::raw_pointer_cast(dvec.data()),data_size,num_of_blks,
		thrust::raw_pointer_cast(ind_vec.data()),num_rounds);


	cudaEventRecord(stop, 0);
	cudaEventSynchronize(stop);
	cudaEventElapsedTime(&elapsedTime, start, stop); 
	exec_time.push_back(elapsedTime);


	cudaEventRecord(start, 0);

	thrust::inclusive_scan(ind_vec.begin(),ind_vec.end(),ind_vec.begin());
	thrust::device_vector<int> output_table(data_size);

	cudaEventRecord(stop, 0);
	cudaEventSynchronize(stop);
	cudaEventElapsedTime(&elapsedTime, start, stop); 
	exec_time.push_back(elapsedTime);


	cudaEventRecord(start, 0);
	//cout << "output hash table" << endl;
	output_hash_table<<<num_of_blks,MAX_HASH_KEY>>>(
		thrust::raw_pointer_cast(dvec.data()),data_size,num_of_blks,
		thrust::raw_pointer_cast(ind_vec.data()),
		thrust::raw_pointer_cast(output_table.data()));

	cudaEventRecord(stop, 0);
	cudaEventSynchronize(stop);
	cudaEventElapsedTime(&elapsedTime, start, stop); 
	exec_time.push_back(elapsedTime);


	cudaEventRecord(start, 0);
	//cout << "sort within blks ---- " << endl;
	sort_within_blks<<<MAX_HASH_KEY,SORTED_SHARED_SIZE>>>(
		thrust::raw_pointer_cast(output_table.data()),
		thrust::raw_pointer_cast(ind_vec.data()),num_of_blks,
		thrust::raw_pointer_cast(dvec.data()),num_rounds);

	cudaEventRecord(stop, 0);
	cudaEventSynchronize(stop);
	cudaEventElapsedTime(&elapsedTime, start, stop); 
	exec_time.push_back(elapsedTime);

	cudaEventRecord(start, 0);
	//cudaDeviceSynchronize();
	thrust::device_vector<int> remove_indicator(data_size,0);
	//cout << "marking duplicates" << endl;

	num_of_blks = (data_size / 1024) + (data_size % 1024 != 0);
	mark_duplicate<<<num_of_blks,1024>>>(
		thrust::raw_pointer_cast(dvec.data()),
		data_size,
		thrust::raw_pointer_cast(remove_indicator.data()));

	cudaEventRecord(stop, 0);
	cudaEventSynchronize(stop);
	cudaEventElapsedTime(&elapsedTime, start, stop); 
	exec_time.push_back(elapsedTime);

	cudaEventRecord(start, 0);
	//cout << "removing" << endl;
	thrust::device_vector<int>::iterator last_elem = thrust::remove_if(dvec.begin(),dvec.begin()+data_size,remove_indicator.begin(),thrust::identity<int>());

	cudaEventRecord(stop, 0);
	cudaEventSynchronize(stop);
	cudaEventElapsedTime(&elapsedTime, start, stop); 
	exec_time.push_back(elapsedTime);

	cudaEventDestroy(start);
	cudaEventDestroy(stop);

	cout << "removing duplicate performance summary" << endl;
	for(int i = 0; i < exec_time.size(); i++){
		cout << "part " << i << " takes time: " << exec_time[i] << endl;
	}

	return (last_elem - dvec.begin());
}

void test_hash_table(){
	srand(1);
	int data_size = 80;
	thrust::host_vector<int> hvec(data_size);
	for(int i = 0; i < data_size; i++) hvec[i] = rand()%data_size;

	vector<vector<int> > hash_table(MAX_HASH_KEY);
	for(int i = 0; i < data_size; i++){
		int value = hvec[i];
		int hash_value = (value * 2654435761) % MAX_HASH_KEY;
		hash_table[hash_value].push_back(value);
	}
	for(int i = 0; i < MAX_HASH_KEY; i++){
		cout << "value " << i << " size " << hash_table[i].size() << ":";
		for(int j = 0; j < hash_table[i].size(); j++){
			int value = hash_table[i][j];
			cout << "(" << value << "," << i << ") | ";
		}
		cout << endl;
	}
	////for(int i = 0; i < data_size; i++){
	////	int value = hvec[i];
	////	int hash_value = (value * 2654435761) % MAX_HASH_KEY;
	////	cout << "(" << value << "," << hash_value << ") | ";
	////}
	cout << endl;
	cout << "computing: " << endl;


	thrust::device_vector<int> dvec = hvec;


	int num_of_blks = (data_size / MAX_HASH_KEY) + (data_size % MAX_HASH_KEY != 0);
	thrust::device_vector<int> ind_vec(num_of_blks * MAX_HASH_KEY);

	/*hash_partition<<<num_of_blks,MAX_HASH_KEY>>>(
		thrust::raw_pointer_cast(dvec.data()),data_size,num_of_blks,
		thrust::raw_pointer_cast(ind_vec.data()));*/

	thrust::inclusive_scan(ind_vec.begin(),ind_vec.end(),ind_vec.begin());
	thrust::device_vector<int> output_table(data_size);

	//cout << "prefix sum: ";
	//for(int i = 0; i < MAX_HASH_KEY; i++){
	//	cout << ind_vec[(i+1)*num_of_blks-1] << " ";
	//}
	//cout << endl;

	output_hash_table<<<num_of_blks,MAX_HASH_KEY>>>(
		thrust::raw_pointer_cast(dvec.data()),data_size,num_of_blks,
		thrust::raw_pointer_cast(ind_vec.data()),
		thrust::raw_pointer_cast(output_table.data()));

	//sort_within_blks<<<MAX_HASH_KEY,SORTED_SHARED_SIZE>>>(
	//	thrust::raw_pointer_cast(output_table.data()),
	//	thrust::raw_pointer_cast(ind_vec.data()),num_of_blks);

	thrust::device_vector<int> remove_indicator(data_size,0);

	num_of_blks = (data_size / 1024) + (data_size % 1024 != 0);
	mark_duplicate<<<num_of_blks,1024>>>(
		thrust::raw_pointer_cast(output_table.data()),
		data_size,
		thrust::raw_pointer_cast(remove_indicator.data()));

	thrust::device_vector<int>::iterator last_elem = thrust::remove_if(output_table.begin(),output_table.end(),remove_indicator.begin(),thrust::identity<int>());
	
	thrust::host_vector<int> result = output_table;

	int iter_key = 0;
	for(int i = 0; i < (last_elem - output_table.begin()); i++){
		int value = result[i];
		int hash_value = (value * 2654435761) % MAX_HASH_KEY;

		if(hash_value != iter_key){
			iter_key = hash_value;
			cout << endl;
		}
		cout << "(" << result[i] << "," << hash_value << ") | ";
	}

}